#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
SRT2EXIF - DJI无人机视频帧提取与GPS信息嵌入工具
===============================================

功能说明:
--------
本工具用于从DJI无人机拍摄的视频中按指定频率提取图像帧，并将SRT字幕文件中的GPS信息
嵌入到对应的图像EXIF数据中。这对于后续进行3D重建、地理信息分析或其他需要地理位置
信息的图像处理非常有用。

用法:
----
python video_extractor.py -s 视频文件路径.MP4 -f 每秒提取帧数 [-m 输出目录] [--log 日志文件]

参数说明:
  -s, --source    必需参数，视频文件路径，程序会自动查找同名的SRT文件
  -f, --fps       必需参数，每秒提取的帧数
  -m, --output    可选参数，输出目录，默认为与视频同名的文件夹
  --log           可选参数，详细日志文件路径

示例:
  python video_extractor.py -s drone_video.MP4 -f 2
  # 从drone_video.MP4中每秒提取2帧，并添加GPS信息

  python video_extractor.py -s drone_video.MP4 -f 5 -m output_frames
  # 从drone_video.MP4中每秒提取5帧，保存到output_frames目录

工作原理:
--------
1. 帧提取逻辑:
   - 读取视频文件，获取其帧率(FPS)
   - 根据用户指定的每秒提取帧数计算跳过帧数
   - 例如: 视频为30FPS，用户要求每秒2帧，则每15帧提取一次

2. SRT解析逻辑:
   - SRT文件包含每一帧的详细元数据，包括GPS坐标、高度、云台角度等
   - 使用状态机解析SRT文件的不同部分(序号、时间戳、字体大小、日期时间、详细信息)
   - 将解析的信息存储在帧信息映射表中，以帧序号为索引

3. EXIF嵌入逻辑:
   - 将GPS坐标转换为EXIF标准的DMS(度分秒)格式
   - 设置经纬度、高度及其参考信息
   - 使用piexif库将信息写入图像文件

注意事项:
--------
- SRT文件必须与视频文件同名，仅扩展名不同
- 确保有足够的磁盘空间存储提取的图像
- 对于大型视频文件，处理可能需要较长时间

作者: 孙雷雷
日期: 2025-04-17
版本: 1.0
"""

import os
import sys
import re
import cv2
import time
import datetime
import argparse
from pathlib import Path
import logging
import piexif
from fractions import Fraction
from tqdm import tqdm
import numpy as np

# 设置日志格式
def setup_logger(log_file=None):
    """
    设置日志记录器，可选择将详细日志写入文件
    """
    # 将控制台日志级别设置为INFO，只显示重要信息
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 移除默认处理器
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)
    
    # 添加控制台处理器
    logging.root.addHandler(console_handler)
    
    # 如果指定了日志文件，添加文件处理器记录详细日志
    if log_file:
        file_handler = logging.FileHandler(log_file)
        file_handler.setLevel(logging.DEBUG)
        file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(file_formatter)
        logging.root.addHandler(file_handler)
    
    # 设置logger的级别为DEBUG，这样所有消息都会被捕获
    logging.root.setLevel(logging.DEBUG)

logger = logging.getLogger(__name__)

def pad_zeros(number, width):
    """
    将数字填充为指定宽度的字符串，不足部分用0填充
    """
    return str(number).zfill(width)

def cv_imread(file_path):
    """
    读取中文路径的图像
    """
    return cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), cv2.IMREAD_COLOR)

def cv_imwrite(file_path, img):
    """
    保存图像到中文路径
    """
    ext = os.path.splitext(file_path)[1]
    cv2.imencode(ext, img)[1].tofile(file_path)

class Exiv2FrameReader:
    """
    处理图像EXIF信息的类
    """
    class FrameTags:
        def __init__(self):
            self.timestamp = "Xmp.exif.REALM.Timestamp"
            self.camera_id = "Exif.Image.Model"
            self.heading = "Xmp.exif.REALM.Heading"
            self.latitude = "Exif.GPSInfo.GPSLatitude"
            self.latituderef = "Exif.GPSInfo.GPSLatitudeRef"
            self.longitude = "Exif.GPSInfo.GPSLongitude"
            self.longituderef = "Exif.GPSInfo.GPSLongitudeRef"
            self.altitude = "Exif.GPSInfo.GPSAltitude"

    def __init__(self):
        self.m_frame_tags = self.FrameTags()

    def add_exiv_info_to_image(self, image_path, latitude, longitude, altitude):
        """
        向图像添加EXIF地理信息
        """
        try:
            # 读取现有EXIF数据
            exif_dict = piexif.load(image_path)
            
            # 确保GPS字典存在
            if "GPS" not in exif_dict:
                exif_dict["GPS"] = {}
            
            # 将经度值转换为DMS格式
            lat_deg = int(abs(latitude))
            lat_min = int((abs(latitude) - lat_deg) * 60)
            lat_sec = int(((abs(latitude) - lat_deg) * 60 - lat_min) * 60 * 10000000)
            
            # 将纬度值转换为DMS格式
            lon_deg = int(abs(longitude))
            lon_min = int((abs(longitude) - lon_deg) * 60)
            lon_sec = int(((abs(longitude) - lon_deg) * 60 - lon_min) * 60 * 10000000)
            
            # 设置经纬度信息
            exif_dict["GPS"][piexif.GPSIFD.GPSLatitude] = (
                (lat_deg, 1),
                (lat_min, 1),
                (lat_sec, 10000000)
            )
            
            exif_dict["GPS"][piexif.GPSIFD.GPSLongitude] = (
                (lon_deg, 1),
                (lon_min, 1),
                (lon_sec, 10000000)
            )
            
            # 设置高度
            if altitude is not None:
                # 直接创建分数表示，避免使用as_integer_ratio
                altitude_int = int(altitude * 100)
                exif_dict["GPS"][piexif.GPSIFD.GPSAltitude] = (altitude_int, 100)
                # 添加高度参考 (0 = 海平面上方)
                exif_dict["GPS"][piexif.GPSIFD.GPSAltitudeRef] = 0
            
            # 设置经纬度参考
            longitude_ref = "E" if longitude >= 0 else "W"
            latitude_ref = "N" if latitude >= 0 else "S"
            exif_dict["GPS"][piexif.GPSIFD.GPSLongitudeRef] = longitude_ref
            exif_dict["GPS"][piexif.GPSIFD.GPSLatitudeRef] = latitude_ref
            
            # 添加GPS版本信息 (2.2.0.0)
            exif_dict["GPS"][piexif.GPSIFD.GPSVersionID] = (2, 2, 0, 0)
            
            # 写入更新的EXIF信息
            exif_bytes = piexif.dump(exif_dict)
            piexif.insert(exif_bytes, image_path)
        except Exception as e:
            logger.error(f"无法添加EXIF数据: {str(e)}")

class FrameInfo:
    """
    存储帧信息的类
    """
    def __init__(self):
        self.rel_timestamp = 0  # 相对时间戳（毫秒）
        self.diff_time = 0      # 时间差（毫秒）
        self.milli_tm = 0       # 毫秒时间
        self.focal_length = 0.0
        self.latitude = 0.0
        self.longitude = 0.0
        self.rel_altitude = 0.0
        self.abs_alt = 0.0
        self.gb_yaw = 0.0
        self.gb_pitch = 0.0
        self.gb_roll = 0.0

class CSRTParser:
    """
    解析SRT字幕文件的类
    """
    # 解析状态机枚举
    BLANK_LINE = 0
    SEQUENCE_NUM = 1
    RELATIVE_TIME = 2
    FONT_SIZE = 3
    DATE_TIME = 4
    DETAIL_INFO = 5
    
    def __init__(self, fpath):
        self.FRAME_CNT = "FrameCnt"
        self.DIFF_TIME = "DiffTime"
        self.FOCAL_LEN = "focal_len"
        self.LATITUDE = "latitude"
        self.LONGITUDE = "longitude"
        self.REL_ALTITUDE = "rel_alt"
        self.ABS_ALTITUDE = "abs_alt"
        self.YAW = "gb_yaw"
        self.PITCH = "gb_pitch"
        self.ROLL = "gb_roll"
        
        self.num2info = {}  # 序列号到帧信息的映射
        
        try:
            with open(fpath, 'r', encoding='utf-8') as file:
                state = self.SEQUENCE_NUM
                seq = 0
                ptr = None
                
                for line in file:
                    line = line.strip()
                    
                    if state == self.BLANK_LINE:
                        state = self.SEQUENCE_NUM
                    
                    elif state == self.SEQUENCE_NUM:
                        ptr = FrameInfo()
                        seq = int(line)
                        self.num2info[seq] = ptr
                        state = self.RELATIVE_TIME
                    
                    elif state == self.RELATIVE_TIME:
                        # 解析时间戳 (00:00:00,000 格式)
                        time_parts = re.match(r'(\d+):(\d+):(\d+),(\d+)', line)
                        if time_parts and ptr:
                            hour = int(time_parts.group(1))
                            minute = int(time_parts.group(2))
                            second = int(time_parts.group(3))
                            millsec = int(time_parts.group(4))
                            ptr.rel_timestamp = (hour * 3600 + minute * 60 + second) * 1000 + millsec
                        state = self.FONT_SIZE
                    
                    elif state == self.FONT_SIZE:
                        if ptr:
                            idx = line.find(self.FRAME_CNT)
                            if idx != -1:
                                # 找到FrameCnt部分
                                frame_cnt_str = line[idx:]
                                # 查找DiffTime部分
                                diff_idx = frame_cnt_str.find(self.DIFF_TIME)
                                if diff_idx != -1:
                                    # 提取DiffTime值，去掉"ms"后缀
                                    diff_time_str = frame_cnt_str[diff_idx + len(self.DIFF_TIME) + 2:]
                                    diff_time_str = diff_time_str.split('ms')[0].strip()
                                    try:
                                        ptr.diff_time = int(diff_time_str)
                                    except ValueError:
                                        logger.warning(f"无法解析DiffTime值: {diff_time_str}")
                            state = self.DATE_TIME
                    
                    elif state == self.DATE_TIME:
                        if ptr:
                            ptr.milli_tm = self.datetime_to_milliseconds(line)
                        state = self.DETAIL_INFO
                    
                    elif state == self.DETAIL_INFO:
                        tokens = self.string_tokenizer(line)
                        if ptr:
                            if self.FOCAL_LEN in tokens:
                                ptr.focal_length = float(tokens[self.FOCAL_LEN])
                            if self.LATITUDE in tokens:
                                ptr.latitude = float(tokens[self.LATITUDE])
                            if self.LONGITUDE in tokens:
                                ptr.longitude = float(tokens[self.LONGITUDE])
                            if self.REL_ALTITUDE in tokens:
                                ptr.rel_altitude = float(tokens[self.REL_ALTITUDE])
                            if self.ABS_ALTITUDE in tokens:
                                ptr.abs_alt = float(tokens[self.ABS_ALTITUDE])
                            if self.YAW in tokens:
                                ptr.gb_yaw = float(tokens[self.YAW])
                            if self.PITCH in tokens:
                                ptr.gb_pitch = float(tokens[self.PITCH])
                            if self.ROLL in tokens:
                                ptr.gb_roll = float(tokens[self.ROLL])
                        state = self.BLANK_LINE
            
            logger.info(f"{len(self.num2info)} records.")
        except Exception as e:
            logger.error(f"Error parsing SRT file: {str(e)}")
    
    def get_frame_info(self, seq):
        """
        获取指定序列号的帧信息
        """
        return self.num2info.get(seq)
    
    def datetime_to_milliseconds(self, datetime_str):
        """
        将日期时间字符串转换为毫秒时间戳
        """
        try:
            # 解析日期时间字符串（格式：YYYY-MM-DD HH:MM:SS.sss）
            dt_parts = datetime_str.split('.')
            dt_obj = None
            ms = 0
            
            if len(dt_parts) > 1:
                dt_obj = datetime.datetime.strptime(dt_parts[0], "%Y-%m-%d %H:%M:%S")
                ms = int(dt_parts[1][:3])  # 取毫秒部分
            else:
                dt_obj = datetime.datetime.strptime(datetime_str, "%Y-%m-%d %H:%M:%S")
            
            # 转换为时间戳（毫秒）
            timestamp = int(dt_obj.timestamp() * 1000) + ms
            return timestamp
        except Exception as e:
            logger.error(f"Failed to parse datetime string: {str(e)}")
            return 0
    
    def string_tokenizer(self, string):
        """
        解析字符串中的键值对
        """
        tokens = {}
        idx1 = string.find('[')
        idx2 = string.find(']')
        tmp = string
        
        while idx1 != -1 and idx2 != -1:
            key = tmp[idx1 + 1:idx2]
            
            # 清理空格和冒号
            key = re.sub(r'  +', ' ', key)
            key = key.replace(' :', ':').replace(': ', ':')
            
            tmp = tmp[idx2 + 1:]
            
            # 解析键值对
            for item in key.split(' '):
                if ':' in item:
                    k, v = item.split(':', 1)
                    tokens[k] = v
            
            idx1 = tmp.find('[')
            idx2 = tmp.find(']')
        
        return tokens

def main():
    """
    主函数
    """
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Extract frames from video and add EXIF data from SRT file')
    parser.add_argument('-s', '--source', required=True, help='Path to the video file (.MP4)')
    parser.add_argument('-f', '--fps', type=int, required=True, help='Number of frames to extract per second')
    parser.add_argument('-m', '--output', help='Output directory for extracted frames (optional)')
    parser.add_argument('--log', help='Path to log file for detailed logging (optional)')
    args = parser.parse_args()
    
    # 设置日志记录
    setup_logger(args.log)
    
    exiv2_reader = Exiv2FrameReader()
    parser = None
    
    video_filepath = args.source
    srt_filepath = video_filepath
    
    # 修改SRT文件路径的扩展名
    idx = srt_filepath.rfind('.')
    if idx != -1:
        srt_filepath = srt_filepath[:idx + 1] + "SRT"
    
    logger.info(f"视频文件: {video_filepath}")
    logger.info(f"SRT文件: {srt_filepath}")
    
    video_path = Path(video_filepath)
    srt_path = Path(srt_filepath)
    
    # 设置输出路径
    if args.output:
        out_image_path = args.output
    else:
        # 如果没有指定输出路径，使用视频文件名作为输出文件夹名
        out_image_path = video_path.parent / video_path.stem
    
    # 检查视频文件是否存在
    if not video_path.exists():
        logger.error(f"视频文件不存在: {video_filepath}")
        return -1
    
    # 检查SRT文件是否存在
    if not srt_path.exists():
        logger.error(f"SRT字幕文件不存在: {srt_filepath}")
    else:
        parser = CSRTParser(srt_filepath)
    
    # 创建输出目录
    out_path = Path(out_image_path)
    if not out_path.exists():
        try:
            out_path.mkdir(parents=True)
        except Exception as e:
            logger.error(f"创建输出目录失败: {out_image_path}, {str(e)}")
            return -1
    
    # 打开视频文件
    video_capture = cv2.VideoCapture(video_filepath)
    if not video_capture.isOpened():
        logger.error(f"无法打开视频文件: {video_filepath}")
        return -1
    
    # 获取视频信息
    fps = video_capture.get(cv2.CAP_PROP_FPS)
    total_frame_count = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
    video_duration = total_frame_count / fps  # 视频总时长（秒）
    
    # 计算帧跳过数量
    # 如果视频是30FPS，用户要求每秒2帧，则需要跳过15帧
    frames_per_second = args.fps
    if frames_per_second <= 0:
        logger.error(f"每秒帧数必须大于0")
        return -1
    
    if frames_per_second > fps:
        logger.warning(f"指定的每秒帧数({frames_per_second})大于视频原始帧率({fps})，将使用原始帧率")
        frames_per_second = fps
    
    frame_skip_num = int(fps / frames_per_second)
    
    logger.info(f"视频帧率 = {fps} FPS, 视频总帧数 = {total_frame_count}, 视频时长 = {video_duration:.2f}秒")
    logger.info(f"每秒抽取 {frames_per_second} 帧, 跳过帧数 = {frame_skip_num}")
    
    frame_number = 0
    extracted_count = 0
    
    # 创建进度条
    pbar = tqdm(total=total_frame_count, desc="抽取视频帧", unit="帧")
    
    while frame_number < total_frame_count:
        # 设置当前帧位置
        video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        retval, frame = video_capture.read()
        
        # 如果读取失败，尝试读取后续帧
        skip_indicator = ""
        for index in range(1, frame_skip_num):
            if not retval:
                skip_indicator += "+"
                video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number + index)
                retval, frame = video_capture.read()
        
        if skip_indicator:
            logger.error(skip_indicator)
        
        if not retval:
            # 更新进度条
            pbar.update(frame_skip_num)
            frame_number += frame_skip_num
            continue
        
        # 保存图像
        pic_standard_number_length = 6
        temp_pic_name = f"img_{pad_zeros(frame_number, pic_standard_number_length)}.jpg"
        output_file = out_path / temp_pic_name
        
        # 使用支持中文路径的方式保存图像
        cv_imwrite(str(output_file), frame)
        extracted_count += 1
        
        # 计算当前时间点（秒）
        current_time = frame_number / fps
        logger.debug(f"[{extracted_count}] 提取帧 {frame_number}/{total_frame_count} (时间: {current_time:.2f}秒) -> {output_file}")
        
        # 添加EXIF信息
        if parser:
            # 获取当前帧对应的SRT帧信息
            # SRT帧从1开始计数，而视频帧从0开始计数
            frame_info = parser.get_frame_info(frame_number + 1)
            if frame_info:
                try:
                    exiv2_reader.add_exiv_info_to_image(
                        str(output_file),
                        frame_info.latitude,
                        frame_info.longitude,
                        frame_info.abs_alt
                    )
                    logger.debug(f"  添加GPS信息: 经度={frame_info.longitude}, 纬度={frame_info.latitude}, 高度={frame_info.abs_alt}")
                except Exception as e:
                    logger.error(f"无法添加EXIF数据到图像 {output_file}: {str(e)}")
            else:
                logger.warning(f"未找到帧 {frame_number + 1} 的信息")
        
        # 更新进度条
        pbar.update(frame_skip_num)
        frame_number += frame_skip_num
    
    # 关闭进度条
    pbar.close()
    
    logger.info(f"完成! 共提取 {extracted_count} 帧图像到 {out_path}")
    return 0

if __name__ == "__main__":
    sys.exit(main())
